#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include "logging.h"
#include <fstream>
#include <opencv2/opencv.hpp>
#include <map>
#include <chrono>

#define TODO()  assert(0)
#define CHECK(status) \
	do\
	{\
		auto ret = (status);\
		if (ret != 0) \
		{\
			std::cerr << "Cuda failure: " << ret << std::endl;\
			abort();\
		}\
	} while(0)

#define DATATYPE DataType::kFLOAT

// input and output blobs
static const int INPUT_H = 32;
static const int INPUT_W = 32;
static const int OUTPUT_SIZE = 10;

const char *INPUT_BLOB_NAME = "data";
const char *OUTPUT_BLOB_NAME = "prob";

unsigned int maxBatchSize = 1;
int batchSize = 1;
const std::string weightsFile = "../lenet5.wts";
const std::string engineFile = "./lenet5.engine";
const std::string test_img_path = "../testImages/mnist_1.jpg";

using namespace nvinfer1;

static Logger gLogger;

inline bool fileExist(const std::string &name)
{
	bool ret;
	std::ifstream f(name.c_str());
	ret = f.good();
	f.close();
	return ret;
}

void printResult(float *res, int size)
{
	for(int i=0; i<size; i++){
		std::cout << res[i] << " ";
	}
	std::cout << std::endl;
}

std::map<std::string, Weights> loadWeights(const std::string file)
{
	std::cout << "Loading weights: " << file << std::endl;
	std::map<std::string, Weights> weightMap;

	std::ifstream input(file);
	assert(input.is_open() && "Unable to load weight file.");

	int32_t count;
	input >> count;
	assert(count > 0 && "Invalid weight map file");

	while(count --)
	{
		Weights wt{DATATYPE, nullptr, 0};
		uint32_t size;

		std::string name;
		input >> name >> std::dec >> size;
		wt.type = DATATYPE;

		uint32_t *val = reinterpret_cast<uint32_t *>(malloc(sizeof(val) * size));
		for(uint32_t x=0, y=size; x < y; ++x)
		{
			input >> std::hex >> val[x];
		}

		wt.values = val;
		wt.count = size;
		weightMap[name] = wt;
	}

	std::cout << "weights file loaded successfully" << std::endl;
	return weightMap;
}

int main(int argc, char *argv[])
{
	/* some vals */
	IBuilder *builder;
	IBuilderConfig *config;
	ICudaEngine *engine;
	INetworkDefinition *network;
	IRuntime *runtime;
	IExecutionContext *context;

	/* PART0: prepare the data */
	float data[INPUT_H * INPUT_W];
	cv::Mat img = cv::imread(test_img_path, cv::IMREAD_GRAYSCALE);
	for(int i=0; i<INPUT_H * INPUT_W; i++){
		data[i] = img.data[i];
	}

	/* PART1: load/build the engine*/
	// if lenet.engine not exists
	if(!fileExist(engineFile)){
		builder = createInferBuilder(gLogger);
		config = builder->createBuilderConfig();

		// load wts to weightMap
		if(!fileExist(weightsFile)){
			std::cerr << "weights files lenet.wts not exists" << std::endl;
			return -1;
		}
		std::map<std::string, Weights> weightMap = loadWeights(weightsFile);
		
		/********* build the network start ********************/
		// network
		network = builder->createNetworkV2(0U);
		ITensor *network_input = network->addInput(INPUT_BLOB_NAME, DATATYPE, Dims3{1, INPUT_W, INPUT_W});
		assert(network_input);

		// conv1
		IConvolutionLayer *conv1 = network->addConvolutionNd(*network_input, 6, DimsHW{5, 5}, weightMap["conv1.weight"], weightMap["conv1.bias"]);
		assert(conv1);
		conv1->setStrideNd(DimsHW{1, 1});

		// relu1
		IActivationLayer *relu1 = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU);
		assert(relu1);

		// pool1
		IPoolingLayer *pool1 = network->addPoolingNd(*relu1->getOutput(0), PoolingType::kAVERAGE, DimsHW{2, 2});
		assert(pool1);
		pool1->setStrideNd(DimsHW{2, 2});

		// conv2
		IConvolutionLayer *conv2 = network->addConvolutionNd(*pool1->getOutput(0), 16, DimsHW{5, 5}, weightMap["conv2.weight"], weightMap["conv2.bias"]);
		assert(conv2);
		conv2->setStrideNd(DimsHW{1, 1});


		// relu2
		IActivationLayer *relu2 = network->addActivation(*conv2->getOutput(0), ActivationType::kRELU);
		assert(relu2);

		// pool2
		IPoolingLayer *pool2 = network->addPoolingNd(*relu2->getOutput(0), PoolingType::kAVERAGE, DimsHW{2, 2});
		assert(pool2);
		pool2->setStrideNd(DimsHW{2, 2});

		// fc1
		IFullyConnectedLayer *fc1 = network->addFullyConnected(*pool2->getOutput(0), 120, weightMap["fc1.weight"], weightMap["fc1.bias"]);
		assert(fc1);

		// relu3
		IActivationLayer *relu3 = network->addActivation(*fc1->getOutput(0), ActivationType::kRELU);
		assert(relu3);

		// fc2
		IFullyConnectedLayer *fc2 = network->addFullyConnected(*relu3->getOutput(0), 84, weightMap["fc2.weight"], weightMap["fc2.bias"]);
		assert(fc2);

		// relu4
		IActivationLayer *relu4 = network->addActivation(*fc2->getOutput(0), ActivationType::kRELU);
		assert(relu4);

		// fc3
		IFullyConnectedLayer *fc3 = network->addFullyConnected(*relu4->getOutput(0), OUTPUT_SIZE, weightMap["fc3.weight"], weightMap["fc3.bias"]);
		assert(fc3);

		// softmax
		ISoftMaxLayer *softmax = network->addSoftMax(*fc3->getOutput(0));
		assert(softmax);
		softmax->getOutput(0)->setName(OUTPUT_BLOB_NAME);
		network->markOutput(*softmax->getOutput(0));
		/********* build the network end *****************/

		// build the engine
		builder->setMaxBatchSize(maxBatchSize);
		config->setMaxWorkspaceSize(1 << 20);
		engine = builder->buildEngineWithConfig(*network, *config);
		assert(engine != nullptr);

		// engine serialize
		IHostMemory *modelStream{nullptr};
		modelStream = engine->serialize();

		std::ofstream p(engineFile.c_str());
		if(!p)
		{
			std::cerr << "could not open plan output file" << std::endl;
			return -1;
		}
		p.write(reinterpret_cast<const char *>(modelStream->data()), modelStream->size());
		modelStream->destroy();

		// destroy the network, engine, builder and release the memory
		network->destroy();
		engine->destroy();
		builder->destroy();

		for(auto &mem : weightMap)
		{
			free((void *)(mem.second.values));
		}
	}

	// load the engine
	char *trtModelStream{nullptr};
	size_t size{0};

	std::ifstream file(engineFile.c_str());
	file.seekg(0, file.end);
	size = file.tellg();
	file.seekg(0, file.beg);
	trtModelStream = new char[size];
	assert(trtModelStream);
	file.read(trtModelStream, size);
	file.close();

	runtime = createInferRuntime(gLogger);
	assert(runtime != nullptr);
	engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr);
	assert(engine != nullptr);
	assert(engine->getNbBindings() == 2);

	/* PART2: do inference */
	float prob[OUTPUT_SIZE];
	void *buffers[2];

	context = engine->createExecutionContext();
	assert(context != nullptr);

	const int inputIndex = engine->getBindingIndex(INPUT_BLOB_NAME);
	const int outputIndex = engine->getBindingIndex(OUTPUT_BLOB_NAME);

	CHECK(cudaMalloc(&buffers[inputIndex], batchSize * INPUT_H * INPUT_W *sizeof(float)));
	CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));

	cudaStream_t stream;
	CHECK(cudaStreamCreate(&stream));

	CHECK(cudaMemcpyAsync(buffers[inputIndex], data, batchSize * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
	context->enqueue(batchSize, buffers, stream, nullptr);
	CHECK(cudaMemcpyAsync(prob, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
	cudaStreamSynchronize(stream);

	cudaStreamDestroy(stream);
	CHECK(cudaFree(buffers[inputIndex]));
	CHECK(cudaFree(buffers[outputIndex]));

	// analyse the result and dispaly img
	printResult(prob, OUTPUT_SIZE);
	cv::imshow("mnist_img", img);
	cv::waitKey(0);

	// destroy the engine 
	context->destroy();
	runtime->destroy();
	engine->destroy();

	return 0;
}
